import torch
from torchvision import models
from efficientnet_pytorch import EfficientNet
import torchxrayvision as xrv
class SimpleCNN(torch.nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()  # b, 3, 32, 32
        layer1 = torch.nn.Sequential()
        layer1.add_module('conv1', torch.nn.Conv2d(3, 32, 3, 1, padding=1))

        # b, 32, 32, 32
        layer1.add_module('relu1', torch.nn.ReLU(True))
        layer1.add_module('pool1', torch.nn.MaxPool2d(2, 2))  # b, 32, 16, 16 //池化为16*16
        self.layer1 = layer1

        layer1.add_module('conv2', torch.nn.Conv2d(32, 64, 3, 1, padding=1))

        # b, 32, 32, 32
        layer1.add_module('relu2', torch.nn.ReLU(True))
        layer1.add_module('pool2', torch.nn.MaxPool2d(2, 2))  # b, 32, 16, 16 //池化为16*16

        layer1.add_module('conv3', torch.nn.Conv2d(64, 128, 3, 1, padding=1))

        # b, 32, 32, 32
        layer1.add_module('relu3', torch.nn.ReLU(True))
        layer1.add_module('pool3', torch.nn.MaxPool2d(2, 2))  # b, 32, 16, 16 //池化为16*16

        layer1.add_module('conv4', torch.nn.Conv2d(128, 256, 3, 1, padding=1))

        # b, 32, 32, 32
        layer1.add_module('relu4', torch.nn.ReLU(True))
        layer1.add_module('pool4', torch.nn.MaxPool2d(2, 2))  # b, 32, 16, 16 //池化为16*16

        layer1.add_module('conv5', torch.nn.Conv2d(256, 256, 3, 1, padding=1))

        # b, 32, 32, 32
        layer1.add_module('relu5', torch.nn.ReLU(True))
        layer1.add_module('pool5', torch.nn.MaxPool2d(2, 2))  # b, 32, 16, 16 //池化为16*16

        ######################

        layer4 = torch.nn.Sequential()
        layer4.add_module('fc1', torch.nn.Linear(12544, 2))
        self.layer4 = layer4

    def forward(self, x):
        conv1 = self.layer1(x)
        fc_input = conv1.view(conv1.size(0), -1)
        fc_out = self.layer4(fc_input)
        return fc_out


model_resnet50 = models.resnet50().cuda()
model_resnet101 = models.resnet101().cuda()
model_resnet152 = models.resnet152().cuda()
model_vgg16 = models.vgg16().cuda()
model_simplecnn = SimpleCNN().cuda()
model_dense121 = models.densenet121().cuda()
model_dense169 = models.densenet169().cuda()
model_efficientnet = EfficientNet.from_name('efficientnet-b5', {'num_classes': 2}).cuda()
model_dense_medical = xrv.models.DenseNet(num_classes=2,in_channels=3).cuda()
# model_inceptionv3 = models.inception_v3().cuda()



'''
每个模型可以分别指定batch size 的大小以及训练的总epoch数量

'''
ModelDict = {
    'resnet50': {'model':model_resnet50,'bs':32,'total_epoch':200} ,
    'resnet101': {'model':model_resnet101,'bs':32,'total_epoch':200} ,
    'resnet152': {'model':model_resnet152,'bs':24,'total_epoch':200} ,
    'vgg16': {'model':model_vgg16,'bs':16,'total_epoch':200} ,
    'simplecnn':{'model':model_simplecnn,'bs':64,'total_epoch':1000},
    'dense121':{'model':model_dense121,'bs':32,'total_epoch':200},
    'dense169':{'model':model_dense169,'bs':32,'total_epoch':200},
    'efficientnet':{'model':model_efficientnet,'bs':16,'total_epoch':200},
    'dense_medical':{'model':model_dense_medical,'bs':32,'total_epoch':200},
    # 'inceptionv3':{'model':model_inceptionv3,'bs':16,'total_epoch':200}
}